import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import joblib
import numpy as np

def train_and_evaluate_salary_models():
    """
    直接从CSV文件读取数据,进行预处理，定义模型字典，训练并评估所有模型，
    并将结果保存到指定路径。
    """
    # 直接在函数内部指定数据文件路径
    file_path = "data/salary.csv"
    # 读取数据
    data = pd.read_csv(file_path)

    salary_data = pd.get_dummies(data, columns=['salary'], drop_first=True)
    salary_data = salary_data.replace(to_replace=" ?", value=np.nan)
    salary_data = salary_data.dropna()

    # 热编码
    salary_data = pd.get_dummies(salary_data, columns=['workclass'])
    salary_data = pd.get_dummies(salary_data, columns=['education'])
    salary_data = pd.get_dummies(salary_data, columns=['marital-status'])
    salary_data = pd.get_dummies(salary_data, columns=['occupation'])
    salary_data = pd.get_dummies(salary_data, columns=['relationship'])
    salary_data = pd.get_dummies(salary_data, columns=['race'])
    salary_data = pd.get_dummies(salary_data, columns=['sex'])
    salary_data = pd.get_dummies(salary_data, columns=['native-country'])

    X, y = salary_data.drop('salary_ >50K', axis=1), salary_data['salary_ >50K']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

    scaler = StandardScaler()
    joblib.dump(X_train.columns, 'model/salary_model/salary_columns.pkl')

    print(X_train.columns)
    X_train = scaler.fit_transform(X_train)
    X_test = scaler.transform(X_test)

    joblib.dump(scaler, 'model/salary_model/salary_scaler.pkl')
    
    # 定义模型字典
    models = {
        "logistic_regression": LogisticRegression(),
        "decision_tree": DecisionTreeClassifier(random_state=42),
        "random_forest": RandomForestClassifier(n_estimators=100, random_state=42),
        "Bayesian": GaussianNB(),
        "SVM": SVC(kernel='linear', probability=True, C=1),
        "KNN": KNeighborsClassifier(n_neighbors=5)
    }
    
    # 训练和评估模型
    for model_name, model in models.items():
        print(f"正在训练和评估 {model_name}...")
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        print(f"{model_name} 准确率: {accuracy}")
        print("分类报告:")
        print(classification_report(y_test, y_pred))
        print("混淆矩阵:")
        print(confusion_matrix(y_test, y_pred))
        # 保存模型
        model_file_path = f'model/salary_model/{model_name}.pkl'
        joblib.dump(model, model_file_path)
        print(f"{model_name}模型已保存")


# 主程序调用
if __name__ == "__main__":
    train_and_evaluate_salary_models()